import json
import pandas as pd

class MemoryHandler:
    def __init__(self, memory_file="memory.json"):
        self.memory_file = memory_file
        self.memory = self.load_memory()

    def load_memory(self):
        """从文件加载记忆数据。"""
        try:
            with open(self.memory_file, 'r') as file:
                return json.load(file)
        except FileNotFoundError:
            return {}

    def save_memory(self):
        """将记忆数据保存到文件。"""
        with open(self.memory_file, 'w') as file:
            json.dump(self.memory, file)

    def update_memory(self, key, value):
        """
        更新记忆中的数据。如果值是 DataFrame、Series 或字典类型，
        将其转换为可以 JSON 序列化的格式。
        """

        def convert_complex_types(val):
            if isinstance(val, pd.DataFrame):
                return val.to_dict(orient='list')
            elif isinstance(val, pd.Series):
                return val.to_dict()
            elif isinstance(val, dict):
                # 递归处理嵌套字典
                return {k: convert_complex_types(v) for k, v in val.items()}
            else:
                return val

        # 转换传入的值，确保它可以被 JSON 序列化
        serializable_value = convert_complex_types(value)

        self.memory[key] = serializable_value
        self.save_memory()



    def get_memory(self, key):
        """检索记忆中的数据。"""
        return self.memory.get(key, None)



